from PIL import Image
import os
import time
import json

import numpy as np
from glob import glob
from tqdm import tqdm

import fire
import sys

sys.setrecursionlimit(1500)

with open('config.json', 'r', encoding='utf-8') as f:
    config = json.load(f)

def sampling():
    # 源样本路径
    sourcePath = config['sourcePath']
    # 目标存放路径
    targetPath = config['targetPath']
    nowtime = time.strftime('%Y%m%d', time.localtime(time.time()))
    # 上一次完成时间
    finishTime = config.get('SamplingFinishTime', "20180101")
    days = os.listdir(sourcePath)
    # 迭代文件
    for filename in days:
        fp = os.path.join(sourcePath, filename)
        filename = os.path.split(fp)[-1]
        if finishTime <= filename <= nowtime and os.path.isdir(fp):
            ## 区域标记 证件
            AreaDetectDatasetsDoc = os.path.join(fp, 'AreaDetect\datasets\DocumentClassifier\datasets')
            sampling_files(AreaDetectDatasetsDoc, sourcePath, targetPath)
            ## 区域标记 版面
            AreaDetectDatasetsLay = os.path.join(fp, 'AreaDetect\datasets\DLayoutClassifier\datasets')
            sampling_files(AreaDetectDatasetsLay, sourcePath, targetPath)
            ## 证件类型采样
            DocumentClassifierDatasets = os.path.join(fp, 'DocumentClassifier\datasets')
            sampling_files(DocumentClassifierDatasets, sourcePath, targetPath)
            ## 版面类型采样
            LayoutClassifierDatasets = os.path.join(fp, 'LayoutClassifier\datasets')
            sampling_files(LayoutClassifierDatasets, sourcePath, targetPath)
            ## 纠偏类型采样
            PreprocessingDatasets = os.path.join(fp, 'Preprocessing')
            sampling_files(PreprocessingDatasets, sourcePath, targetPath)
            ## 图片类型采样
            TypeClassifierDatasets = os.path.join(fp, 'TypeClassifier\datasets')
            sampling_files(TypeClassifierDatasets, sourcePath, targetPath)
            ## 图片类型采样
            RecallDatasets = os.path.join(fp, 'Recall\datasets')
            recall_sampling_files(RecallDatasets, sourcePath, targetPath)

            config['SamplingFinishTime'] = filename
    # 采样完成后记录采样时间
    f.close()
    with open('config.json', 'w', encoding='utf-8') as outfile:
        outfile.write(json.dumps(config, sort_keys=False, indent=2, separators=(',', ':'), ensure_ascii=False))
        outfile.close()
    pass

def recall_sampling_files(RecallDatasets, sourcePath, targetPath):
    if os.path.exists(RecallDatasets):
        for filename_ in os.listdir(RecallDatasets):
            fp_ = os.path.join(RecallDatasets, filename_,'0')
            datasetsFiles = glob(os.path.join(fp_, '*'))
            for i in tqdm(range(len(datasetsFiles)), desc=fp_):
                src = datasetsFiles[i]
                target = src.replace(sourcePath, targetPath)
                sampling_file(src, target)
    pass

def sampling_files(datasets, sourcePath, targetPath):
    if os.path.exists(datasets) :
        for filename_ in os.listdir(datasets):
            fp_ = os.path.join(datasets, filename_)
            datasetsFiles = glob(os.path.join(fp_, '*'))
            for i in tqdm(range(len(datasetsFiles)), desc=fp_):
                src = datasetsFiles[i]
                target = src.replace(sourcePath, targetPath)
                sampling_file(src, target)
    pass

def copy_file(src, target):
    filename = os.path.split(src)[-1]
    if not os.path.exists(target.replace(filename, '')):
        os.makedirs(target.replace(filename, ''))
    if not os.path.exists(target) and os.path.isfile(src):
        image1 = Image.open(src)
        image1.save(target)
    pass

def sampling_file(src, target):
    filename = os.path.split(src)[-1]
    if not os.path.exists(target.replace(filename, '')):
        os.makedirs(target.replace(filename, ''))
    if not os.path.exists(target):
        image2 = resize2(Image.open(src), (256, 256))
        image2.save(target)
    pass

def resize2(img, size):
    image = img
    fg_w, fg_h = image.size
    bg_w, bg_h = size

    if image.mode == 'L':
        out_image = np.zeros((bg_h, bg_w), dtype='uint8') + 255
    elif image.mode == 'RGB':
        out_image = np.zeros((bg_h, bg_w, 3), dtype='uint8') + 255
    elif image.mode == 'RGBA':
        out_image = np.zeros((bg_h, bg_w, 4), dtype='uint8') + 255
    else:
        image = image.convert('RGB')
        out_image = np.zeros((bg_h, bg_w, 3), dtype='uint8') + 255

    if fg_h == fg_w:
        image2 = image.resize(size, resample=Image.BOX)
        out_image = np.array(image2, dtype='uint8')
    elif fg_h > fg_w:
        h = bg_h
        w = int(h / fg_h * fg_w)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[:, (bg_w - w) // 2: (bg_w - w) // 2 + w] = image2_arr
    elif fg_w > fg_h:
        w = bg_w
        h = int(w / fg_w * fg_h)
        image2 = image.resize((w, h), resample=Image.BOX)
        image2_arr = np.array(image2, dtype='uint8')
        out_image[(bg_h - h) // 2: (bg_h - h) // 2 + h, :] = image2_arr
    return Image.fromarray(out_image)

if __name__ == '__main__':
    fire.Fire(sampling)
